# Copyright 2021 DeepMind Technologies Limited.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Organic molecule config. Bicyclobutane, butadiene and cyclobutadiene."""
# Geometries for bicyclobutane to butadiene transition taken from
# CASSCF(10,10)/cc-pVDZ calculations used in
# A. Kinal and P. Piecuch, J. Phys. Chem. A 2007, 111, 734-742
# Geometries for cyclobutadiene automerization taken from MR-BWCCSD(T)
# a.c./cc-pVTZ calculations from
# Bhaskaran-Nair et al., J. Chem. Phys. 2008, 129, 184104, and are the same as
# those in J. Hermann, Z. Schätzle and F Noé, Nat. Chem. 2020, 12, 891–897

from lapnet import base_config
from lapnet.utils import system

systems = {
    'bicbut': [['C', (1.0487346562, 0.5208579773, 0.2375867187)],
               ['C', (0.2497284256, -0.7666691493, 0.0936474818)],
               ['C', (-0.1817326465, 0.4922777820, -0.6579637266)],
               ['C', (-1.1430708301, -0.1901383337, 0.3048494250)],
               ['H', (2.0107137141, 0.5520589541, -0.2623459977)],
               ['H', (1.0071921280, 1.0672669240, 1.1766131856)],
               ['H', (0.5438033167, -1.7129829738, -0.3260782874)],
               ['H', (-0.2580605320, 0.6268443026, -1.7229636111)],
               ['H', (-1.3778676954, 0.2935640723, 1.2498189977)],
               ['H', (-1.9664163102, -0.7380906148, -0.1402911727)]],
    'con_TS': [['C', (1.0422528085, 0.5189448459, 0.2893513723)],
               ['C', (0.6334392052, -0.8563584473, -0.1382423606)],
               ['C', (-0.2492035181, 0.3134656784, -0.5658962512)],
               ['C', (-1.3903646889, 0.0535204487, 0.2987506023)],
               ['H', (1.8587636947, 0.9382817031, -0.2871146890)],
               ['H', (0.9494853889, 0.8960565051, 1.3038563129)],
               ['H', (0.3506375894, -1.7147937260, 0.4585707483)],
               ['H', (-0.3391417369, 0.6603641863, -1.5850373819)],
               ['H', (-1.2605467656, 0.0656225945, 1.3701508857)],
               ['H', (-2.3153892612, -0.3457478660, -0.0991685880)]],
    'dis_TS': [['C', (1.5864390444, -0.1568990400, -0.1998155990)],
               ['C', (-0.8207390911, 0.8031532550, -0.2771554962)],
               ['C', (0.2514913592, 0.0515423448, 0.4758741643)],
               ['C', (-1.0037104567, -0.6789877402, -0.0965401189)],
               ['H', (2.4861305372, 0.1949133826, 0.2874101433)],
               ['H', (1.6111805503, -0.2769458302, -1.2753251100)],
               ['H', (-1.4350764228, 1.6366792379, 0.0289087336)],
               ['H', (0.2833919284, 0.1769734467, 1.5525271253)],
               ['H', (-1.7484283536, -1.0231589431, 0.6120702030)],
               ['H', (-0.8524391649, -1.3241689195, -0.9544331346)]],
    'g-but': [['C', (1.4852019019, 0.4107781008, 0.5915178362)],
              ['C', (0.7841417614, -0.4218449588, -0.2276848579)],
              ['C', (-0.6577970182, -0.2577617373, -0.6080850660)],
              ['C', (-1.6247236649, 0.2933006709, 0.1775352473)],
              ['H', (1.0376813593, 1.2956518484, 1.0267024109)],
              ['H', (2.5232360753, 0.2129135014, 0.8248568552)],
              ['H', (1.2972328960, -1.2700686671, -0.6686116041)],
              ['H', (-0.9356614935, -0.6338686329, -1.5871170536)],
              ['H', (-1.4152018269, 0.6472889925, 1.1792563311)],
              ['H', (-2.6423222755, 0.3847635835, -0.1791755263)]],
    'gt-TS': [['C', (1.7836595975, 0.4683155866, -0.4860478101)],
              ['C', (0.7828892933, -0.4014025715, -0.1873880949)],
              ['C', (-0.6557274850, -0.2156646805, -0.6243545354)],
              ['C', (-1.6396999531, 0.2526943506, 0.1877948644)],
              ['H', (1.6003117673, 1.3693309737, -1.0595471944)],
              ['H', (2.7986234673, 0.2854595500, -0.1564989895)],
              ['H', (1.0128486304, -1.2934621995, 0.3872559845)],
              ['H', (-0.9003245968, -0.4891235826, -1.6462438855)],
              ['H', (-1.4414954784, 0.5345813494, 1.2152198579)],
              ['H', (-2.6556262424, 0.3594422237, -0.1709361970)]],
    't-but': [['C', (0.6109149108, 1.7798412991, -0.0000000370)],
              ['C', (0.6162339625, 0.4163908910, -0.0000000070)],
              ['C', (-0.6162376752, -0.4163867945, -0.0000000601)],
              ['C', (-0.6109129465, -1.7798435851, 0.0000000007)],
              ['H', (1.5340442204, 2.3439205382, 0.0000000490)],
              ['H', (-0.3156117962, 2.3419017314, 0.0000000338)],
              ['H', (1.5642720455, -0.1114324578, -0.0000000088)],
              ['H', (-1.5642719469, 0.1114307897, -0.0000000331)],
              ['H', (-1.5340441021, -2.3439203971, 0.0000000714)],
              ['H', (0.3156133277, -2.3419020150, -0.0000000088)]],
    'cycbut-ground': [['C', (0.0000000e+00, 0.0000000e+00, 0.0000000e+00)],
                      ['C', (2.9555318e+00, 0.0000000e+00, 0.0000000e+00)],
                      ['C', (2.9555318e+00, 2.5586891e+00, 0.0000000e+00)],
                      ['C', (0.0000000e+00, 2.5586891e+00, 0.0000000e+00)],
                      ['H', (-1.4402903e+00, -1.4433100e+00, 1.7675451e-16)],
                      ['H', (4.3958220e+00, -1.4433100e+00, -1.7675451e-16)],
                      ['H', (4.3958220e+00, 4.0019994e+00, 1.7675451e-16)],
                      ['H', (-1.4402903e+00, 4.0019994e+00, -1.7675451e-16)]],
    'cycbut-trans': [['C', (0.0000000e+00, 0.0000000e+00, 0.0000000e+00)],
                     ['C', (2.7419927e+00, 0.0000000e+00, 0.0000000e+00)],
                     ['C', (2.7419927e+00, 2.7419927e+00, 0.0000000e+00)],
                     ['C', (0.0000000e+00, 2.7419927e+00, 0.0000000e+00)],
                     ['H', (-1.4404647e+00, -1.4404647e+00, 1.7640606e-16)],
                     ['H', (4.1824574e+00, -1.4404647e+00, -1.7640606e-16)],
                     ['H', (4.1824574e+00, 4.1824574e+00, 1.7640606e-16)],
                     ['H', (-1.4404647e+00, 4.1824574e+00, -1.7640606e-16)]]
}


def organic_molecule(cfg):
  """Sets molecule and electrons fields in ConfigDict."""
  cfg.system.electrons = (15, 15)
  units = 'angstrom'
  if 'cycbut' in cfg.system.molecule_name:
    cfg.system.electrons = (14, 14)
    units = 'bohr'
  if cfg.system.molecule_name not in systems.keys():
    raise ValueError(f'Unrecognized molecule: {cfg.system.molecule_name}')
  molecule = []
  for element, coords in systems[cfg.system.molecule_name]:
    molecule.append(system.Atom(symbol=element, coords=coords, units=units))
  cfg.system.molecule = molecule
  return cfg


def get_config():
  """Returns config for running bibut."""
  name = 'organic'
  cfg = base_config.default()
  cfg.system.molecule_name = name
  with cfg.ignore_type():
    cfg.system.set_molecule = organic_molecule
    cfg.config_module = '.organic'
  return cfg
